package com.lamb.script.core.executor.loader;

import com.lamb.script.core.utils.ScriptConsole;
import com.lamb.script.core.utils.ScriptUtil;

import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Map;
import java.util.Vector;
import java.util.concurrent.ConcurrentHashMap;
import java.util.jar.JarFile;
import java.util.jar.Manifest;

/**
 * 自定义类加载器.
 *
 * @author : lamb
 * @version : 1.0
 * @date :  2019-08-02 14:31
 **/
public class ScriptClassLoader extends URLClassLoader {

    private static final String DEFAULT_LOCAL_REPOSITORY = "/tmp/.script.lib";
    private static final String DEFAULT_MIRROR_URL = "http://maven.aliyun.com/nexus/content/groups/public";
    private static final String SPLIT_FILE_PATH = "/";
    private static final String SPLIT_SPOT = ".";
    private static final String PROTOCOL_SPIT = "!/";

    private final Map<URL, byte[]> classCache = new ConcurrentHashMap<>();
    private final Map<String, Manifest> manifestCache = new ConcurrentHashMap<>();
    private final HashSet<String> jarNameCache = new HashSet<>();
    /**
     * url处理器.
     */
    private final ScriptUrlStreamHandler scriptURLStreamHandler;
    /**
     * 脚本输出器.
     */
    private final ScriptConsole scriptConsole;

    private MavenJarLoader mavenJarLoader;

    public ScriptClassLoader(URL[] urls, ClassLoader contextClassLoader, ScriptConsole scriptConsole) {
        super(urls, contextClassLoader);
        this.scriptConsole = scriptConsole;
        this.scriptURLStreamHandler = new ScriptUrlStreamHandler(this);
        // 支持可配置.
        this.mavenJarLoader = new MavenJarLoader(
                System.getProperty("localRepository", DEFAULT_LOCAL_REPOSITORY),
                System.getProperty("mirrorUrl", DEFAULT_MIRROR_URL),
                scriptConsole);
    }

    @Override
    protected Class<?> findClass(String className) throws ClassNotFoundException {
        String classKey = className.replace(SPLIT_SPOT, SPLIT_FILE_PATH) + ".class";
        URL resource = findResource(classKey);
        if (null == resource) {
            throw new ClassNotFoundException(className);
        }

        byte[] bytes = classCache.get(resource);
        if (className.contains(SPLIT_SPOT)) {
            String packageName = className.substring(0, className.lastIndexOf(SPLIT_SPOT));
            Package aPackage = super.getPackage(packageName);
            String javaName = resource.toString();
            javaName = javaName.substring(0, javaName.indexOf(PROTOCOL_SPIT) + PROTOCOL_SPIT.length());
            Manifest manifest = manifestCache.get(javaName);
            // 未加载过则忽略.
            if (null == aPackage && null != manifest) {
                String spec = javaName + packageName.replace(SPLIT_SPOT, SPLIT_FILE_PATH);
                URL url;
                try {
                    url = new URL(null, spec, scriptURLStreamHandler);
                } catch (MalformedURLException e) {
                    throw new ClassNotFoundException("error:" + spec, e);
                }
                super.definePackage(packageName, manifest, url);
            }
        }
        return super.defineClass(className, bytes, 0, bytes.length);
    }

    @Override
    public URL findResource(String name) {
        Enumeration<URL> resources = findResources(name);
        if (resources.hasMoreElements()) {
            return resources.nextElement();
        }
        return null;
    }

    @Override
    public Enumeration<URL> findResources(String name) {
        Vector<URL> result = new Vector<>();
        name = PROTOCOL_SPIT + name;
        for (URL url : classCache.keySet()) {
            if (url.getPath().endsWith(name)) {
                result.add(url);
            }
        }
        return result.elements();
    }

    /**
     * 添加jar包.
     *
     * @param filePath jar地址.
     */
    public void addJar(String filePath) throws IOException {
        if (!jarNameCache.add(filePath)) {
            return;
        }
        File file = mavenJarLoader.findJar(filePath);
        JarFile jarFile = new JarFile(file);
        final String baseUrl = "jar:" + file.getName() + PROTOCOL_SPIT;
        Manifest manifest = jarFile.getManifest();
        manifestCache.put(baseUrl, manifest);
        jarFile.stream().forEach(jarEntry -> {
            try {
                String name = jarEntry.getName();
                URL url = new URL(null, baseUrl + name, scriptURLStreamHandler);

                // 文件夹忽略.
                if (name.endsWith(SPLIT_FILE_PATH)) {
                    classCache.put(url, new byte[0]);
                    if (name.startsWith("META-INF/")) {
                        return;
                    }
                    return;
                }
                classCache.put(url, ScriptUtil.readToByteBuffer(jarFile.getInputStream(jarEntry)));
            } catch (Exception e) {
                scriptConsole.error("addJar error, filePath=" + filePath);
                e.printStackTrace(scriptConsole.err);
            }
        });
    }

    /**
     * 获取内容.
     *
     * @param url 地址.
     * @return 内容.
     */
    public byte[] getContentByUrl(URL url) {
        return classCache.get(url);
    }
}
